import json
from fastapi import APIRouter, HTTPException

from db.Do import BaseReq, we_library, ChatHistory, ChatRole
from data import use_llm
import math
import ast

router = APIRouter()

from PIL import Image
import requests
from io import BytesIO
import base64


import numpy as np
from sentence_transformers import SentenceTransformer
import faiss
import json


class LongTermMemory:
    def __init__(self, embedding_dim=384):
        self.encoder = SentenceTransformer('all-MiniLM-L6-v2')
        self.index = faiss.IndexFlatL2(embedding_dim)
        self.memory_store = []

    def add_memory(self, text: str, metadata: dict = None):
        """添加新记忆到长期存储"""
        embedding = self.encoder.encode([text])[0]
        self.index.add(np.array([embedding]))
        self.memory_store.append({
            "text": text,
            "embedding": embedding.tolist(),
            "metadata": metadata or {}
        })

    def retrieve(self, query: str, k=5) -> list:
        """检索相关记忆"""
        query_embed = self.encoder.encode([query])[0]
        distances, indices = self.index.search(np.array([query_embed]), k)

        results = []
        for i in indices[0]:
            if i >= 0:  # FAISS可能返回-1
                results.append(self.memory_store[i])
        return results

    def summarize_conversation(self, messages: list) -> str:
        """生成对话摘要"""
        # 简化的摘要生成，实际应用中可以使用大模型生成
        key_points = []
        for msg in messages:
            if msg['role'] == 'user':
                key_points.append(f"用户: {msg['content']}")
            elif msg['role'] == 'assistant':
                key_points.append(f"助手: {msg['content']}")
        return "\n".join(key_points[-3:])  # 取最近3条作为摘要





class ImageProcessor:
    def __init__(self):
        # 实际应用中应使用多模态模型
        self.model = None  # 预留多模态模型接口

    def describe_image(self, image_data: str) -> str:
        """处理并描述图片内容"""
        try:
            # 处理base64编码的图片
            if image_data.startswith("data:image"):
                header, encoded = image_data.split(",", 1)
                image_bytes = base64.b64decode(encoded)
            # 处理URL
            elif image_data.startswith("http"):
                response = requests.get(image_data)
                image_bytes = response.content
            else:
                return "无法识别的图片格式"

            image = Image.open(BytesIO(image_bytes))

            # 简化处理 - 实际应使用多模态模型
            # 示例: return self.model.describe(image)
            return f"图片描述: {image.format}格式, 尺寸{image.width}x{image.height}"

        except Exception as e:
            return f"图片处理错误: {str(e)}"


class ToolCallingSystem:
    def __init__(self):
        self.tools = {
            "calculator": {
                "function": self.calculate,
                "description": "执行数学计算",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "expression": {
                            "type": "string",
                            "description": "需要计算的数学表达式，例如：'3 + 5 * 2'"
                        }
                    },
                    "required": ["expression"]
                }
            },
            "web_search": {
                "function": self.web_search,
                "description": "执行网络搜索",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": "搜索关键词"
                        }
                    },
                    "required": ["query"]
                }
            },
            "code_executor": {
                "function": self.execute_code,
                "description": "在安全环境中执行代码",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "language": {
                            "type": "string",
                            "enum": ["python", "javascript", "bash"],
                            "description": "编程语言"
                        },
                        "code": {
                            "type": "string",
                            "description": "需要执行的代码"
                        }
                    },
                    "required": ["language", "code"]
                }
            },
            "calendar": {
                "function": self.calendar_event,
                "description": "创建日历事件",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "title": {
                            "type": "string",
                            "description": "事件标题"
                        },
                        "date": {
                            "type": "string",
                            "description": "事件日期和时间，格式：YYYY-MM-DD HH:MM"
                        }
                    },
                    "required": ["title", "date"]
                }
            }
        }

    def get_tool_definitions(self) -> list:
        """获取工具定义列表（符合OpenAI函数调用规范）"""
        definitions = []
        for name, tool in self.tools.items():
            definitions.append({
                "type": "function",
                "function": {
                    "name": name,
                    "description": tool["description"],
                    "parameters": tool["parameters"]
                }
            })
        return definitions

    def execute_tool(self, tool_name: str, parameters: dict) -> str:
        """执行指定工具"""
        if tool_name not in self.tools:
            return f"错误: 工具 '{tool_name}' 不存在"

        try:
            return self.tools[tool_name]["function"](**parameters)
        except Exception as e:
            return f"工具执行错误: {str(e)}"

    def calculate(self, expression: str) -> str:
        """计算数学表达式"""
        try:
            # 安全计算 - 限制可用函数
            allowed_names = {}
            safe_dict = {k: v for k, v in math.__dict__.items() if not k.startswith("__")}
            safe_dict.update(allowed_names)

            # 使用ast解析更安全
            node = ast.parse(expression, mode='eval')
            if not all(
                    isinstance(n, (ast.Expression, ast.Num, ast.BinOp, ast.UnaryOp, ast.Name)) for n in ast.walk(node)):
                return "表达式包含不安全操作"

            return str(eval(compile(node, '<string>', 'eval'), {"__builtins__": {}}, safe_dict))
        except Exception as e:
            return f"计算错误: {str(e)}"

    def web_search(self, query: str) -> str:
        """执行网络搜索"""
        # 实际应用中应调用搜索引擎API
        try:
            # 模拟搜索 - 实际应用中替换为API调用
            results = [
                {"title": f"关于{query}的结果1", "url": "https://example.com/result1"},
                {"title": f"关于{query}的结果2", "url": "https://example.com/result2"},
            ]
            return json.dumps(results, ensure_ascii=False)
        except Exception as e:
            return f"搜索错误: {str(e)}"

    def execute_code(self, language: str, code: str) -> str:
        """在安全环境中执行代码"""
        # 实际应用中应使用沙箱环境
        if language == "python":
            try:
                # 安全执行 - 使用受限环境
                restricted_globals = {"__builtins__": None}
                restricted_locals = {}

                # 使用exec执行代码
                exec(code, restricted_globals, restricted_locals)

                # 返回结果（如果有）
                result = restricted_locals.get('result', '代码执行完成')
                return str(result)
            except Exception as e:
                return f"代码执行错误: {str(e)}"
        else:
            return f"不支持的语言: {language}"

    def calendar_event(self, title: str, date: str) -> str:
        """创建日历事件"""
        # 模拟日历集成
        try:
            # 实际应用中应调用日历API
            event_id = f"event_{hash(title + date)}"
            return f"已创建事件: '{title}' 于 {date} (ID: {event_id})"
        except Exception as e:
            return f"创建事件错误: {str(e)}"


















def get_role_system_message(role_id: str) -> dict:
    """获取角色系统消息"""
    if not role_id or role_id == "请选择角色":
        return None

    role_info = we_library.fetch_one(
        "SELECT id, role_name, role_setting FROM chat_role WHERE id = ?;",
        (role_id,)
    )

    if not role_info:
        raise HTTPException(status_code=404, detail="角色未找到")

    return {
        "role": "system",
        "content": f"名称：{role_info['role_name']}，描述：{role_info['role_setting']}"
    }


def load_chat_history(chat_id: int) -> list:
    """加载聊天历史"""
    if not chat_id or chat_id == 0:
        return []

    db_record = we_library.fetch_one(
        "SELECT content FROM chat_history WHERE id = ?;",
        (chat_id,)
    )

    return json.loads(db_record["content"]) if db_record and db_record.get("content") else []


def save_chat_history(chat_id: int, messages: list, prompt: str) -> int:
    """保存聊天记录到数据库"""
    # 创建ChatHistory实例
    chat_history = ChatHistory(
        table_name="chat_history",
        introduce=prompt if chat_id == 0 else None,
        id=chat_id if chat_id != 0 else None,
        content=json.dumps(messages, ensure_ascii=False)
    )

    # 保存到数据库并返回新ID
    return we_library.add_or_update(chat_history, "chat_history")


def enhanced_generate_response(messages: list, tools: list = None) -> dict:
    """
    增强版大模型交互，支持工具调用
    返回结构：{"type": "text" | "tool_calls", "content": str | list}
    """
    # 如果有工具定义，添加工具描述到系统消息
    if tools:
        tool_system_message = {
            "role": "system",
            "content": "你可以使用以下工具解决问题:",
            "tools": tools
        }
        # 确保系统消息在最前面
        if messages and messages[0]['role'] == 'system':
            messages[0] = tool_system_message
        else:
            messages.insert(0, tool_system_message)

    # 调用大模型
    raw_response = use_llm._generate_response(messages)

    # 尝试解析工具调用响应（如果返回的是JSON字符串）
    try:
        # 尝试解析为JSON
        parsed_response = json.loads(raw_response)

        # 检查是否是工具调用格式
        if isinstance(parsed_response, dict) and "tool_calls" in parsed_response:
            return {
                "type": "tool_calls",
                "content": parsed_response["tool_calls"]
            }
    except json.JSONDecodeError:
        # 不是JSON格式，继续处理
        pass

    # 尝试检测自定义工具调用格式
    if isinstance(raw_response, str) and "TOOL_CALL:" in raw_response:
        return parse_custom_tool_calls(raw_response)

    # 普通文本响应
    return {
        "type": "text",
        "content": raw_response
    }


def parse_custom_tool_calls(response_str: str) -> dict:
    """
    解析自定义工具调用格式
    示例格式:
    TOOL_CALL: {"name": "calculator", "arguments": {"expression": "2+2"}}
    """
    tool_calls = []

    # 分割多个工具调用
    call_blocks = response_str.split("TOOL_CALL:")

    for block in call_blocks:
        block = block.strip()
        if not block:
            continue

        try:
            # 尝试解析JSON
            tool_call = json.loads(block)

            # 验证基本结构
            if not isinstance(tool_call, dict):
                continue

            if "name" not in tool_call or "arguments" not in tool_call:
                continue

            tool_calls.append({
                "id": f"call_{len(tool_calls)}",
                "type": "function",
                "function": {
                    "name": tool_call["name"],
                    "arguments": json.dumps(tool_call["arguments"], ensure_ascii=False)
                }
            })
        except json.JSONDecodeError:
            # 不是有效的JSON，跳过
            continue

    if tool_calls:
        return {
            "type": "tool_calls",
            "content": tool_calls
        }

    # 如果没有找到有效工具调用，返回原始文本
    return {
        "type": "text",
        "content": response_str
    }






# if __name__ == '__main__':
#     # 创建新对话
#     response = chat(BaseReq(
#         prompt="计算圆的面积，半径是5cm",
#         image_data=None
#     ))
#
#     # 带图片的对话
#     response = chat(BaseReq(
#         prompt="描述这张图片的内容",
#         image_data="..."
#     ))
#
#     # 工具调用
#     response = chat(BaseReq(
#         prompt="请搜索最近的AI技术新闻"
#     ))
#
#     # 代码执行
#     response = chat(BaseReq(
#         prompt="请写一个Python函数计算斐波那契数列并执行"
#     ))